import numpy as np


class kmeans:
    def __init__(self, k, maxiter, data, labels):
        self.k = k
        self.maxiter = maxiter
        self.data = data
        # 由图片28*28二维矩阵转化为784一维矩阵
        self.labels = labels
        self.distances = np.zeros((self.data.shape[0], self.k))
        self.center = np.zeros((self.k, self.data.shape[1]))

    def get_distances(self):
        for i in range(self.data.shape[0]):
            distance_i = ((np.tile(self.data[i], (self.k, 1)) - self.center) ** 2).sum(axis=1) ** 0.5
            self.distances[i] = distance_i

    def get_center(self):
        self.classifications = np.random.randint(0, self.k, (self.data.shape[0]))
        for i in range(self.k):
            self.classifications[i] = i

    def classify(self):
        new_classifications = np.argmin(self.distances, axis=1)
        if any(self.classifications - new_classifications):
            self.classifications = new_classifications
            return True
        else:
            return False

    def update_center(self):
        for i in range(self.k):
            self.center[i] = np.mean(self.data[self.classifications == i], axis=0)

    def fit(self):
        self.get_center()
        for i in range(self.maxiter):
            self.update_center()
            self.get_distances()
            if not self.classify():
                break
